# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-08-29

from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout

from Common.config_utils import Config


class LSTMModel:
    """
    LSTM模型，用于处理时间序列预测任务。
    :param input_dim: int，输入数据的维度。
    :param output_dim: int，预测目标的维度。
    """
    def __init__(
            self,
            input_dim,
            output_dim,
    ):
        """
        初始化 LSTM 模型。
        :param input_dim: int，输入数据的维度。
        :param output_dim: int，预测目标的维度。
        """
        self.input_shape = input_dim
        self.output_shape = output_dim
        self.model = self.build_model()

    def build_model(
            self,
    ):
        """
        使用全局配置中的超参数构建 LSTM 模型。
        :return: 构建好的 LSTM 模型。
        """
        # 从全局配置中导入超参数
        num_layers = Config.NUM_LAYERS
        units = Config.UNITS
        dropout_rate = Config.DROPOUT_RATE

        # 初始化 LSTM 网络
        model = Sequential()
        model.add(LSTM(units, input_shape=self.input_shape, return_sequences=True))
        for _ in range(1, num_layers):
            model.add(LSTM(units, return_sequences=False if _ == num_layers - 1 else True))
        model.add(Dropout(dropout_rate))
        model.add(Dense(self.output_shape))
        model.compile(optimizer='adam', loss='mse', metrics=['mse'])
        return model

    def train(
            self,
            x_train,
            y_train,
    ):
        """
        训练模型。
        :param x_train: numpy.ndarray，训练数据的特征。
        :param y_train: numpy.ndarray，训练数据的目标值。
        """
        epochs = Config.EPOCHS
        batch_size = Config.BATCH_SIZE
        self.model.fit(x_train, y_train, epochs=epochs, batch_size=batch_size)

    def predict(
            self,
            x,
    ):
        """
        使用模型进行预测。
        :param x: numpy.ndarray，输入数据的特征。
        :return: numpy.ndarray，模型预测的结果。
        """
        return self.model.predict(x)
